# Copyright (C) 2007 Richard Boulton
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
r"""requestqueue.py: A threadsafe request queue implementation.

The queue allows requests to be appended, with an associated priority.  Workers
can poll or wait for the queue to pass requests to them.  Workers can then
return results for the requests, which the original callers can pick up using a
request ID or by blocking.

"""

from bisect import bisect_left
import threading
import time

__all__ = (
    'RequestQueue',
    'QueueException',
    'ShutdownException',
)

class QueueException(Exception):
    """Base class for all exceptions raised by a Queue.

    """
    pass

class ShutdownException(QueueException):
    """An execption thrown when a task is added to a queue after its shutdown.

    """
    pass

class Request(object):
    __slots__ = ('data', 'priority', 'noreturn', 'id', 'queue')
    def __init__(self, data, priority, noreturn, id, queue):
        """Make a new Request object.

        "queue" is the queue the Request object came from.

        "data" is the user-specified data describing the request.

        "priority" contains the priority of the request to insert on the queue.

        "noreturn" is a boolean - if True, the request cannot return results.

        "id" is the ID number for the request

        """
        self.data = data
        self._priority = priority
        self._noreturn = noreturn
        self._id = id
        self._queue = queue

    def finished(self, result):
        """Notify the queue that this request has finished.

        "result" is the result of the request.

        This must be called exactly once, to indicate that each request has
        been processed.

        """
        if self._noreturn:
            return
        if self.queue is None:
            raise QueueException("Request.finished() was called more than once")
        self.queue._finished(self._id, result)
        self.queue = None

class RequestResult(object):
    """An object encapsulating the result of a request.

    The id of the request is available in "id", and the result of the request
    is in "result".

    """
    __slots__ = ('id', 'result')
    def __init__(self, id, result):
        self.id = id
        self.result = result

def _wait_until_timeout(condition, endtime):
    """Wait for a condition to be signalled, until endtime.

    If endtime is None, wait indefinitely.

    The condition must be acquired before this is called.

    Returns True if the time has passed endtime, False otherwise.

    """
    if endtime is None:
        condition.wait()
    else:
        waittime = endtime - time.time()
        if waittime <= 0:
            return True
        condition.wait(waittime)
    return False

class RequestQueue(object):
    def __init__(self):
        """Create a new request queue.

        The queue allows requests to be placed on the queue, with a given
        priority.  Requests can then be picked from the queue (perhaps by a
        separate thread), processed, and the results can then be returned to
        the queue.  Results for a request can be actively polled for or waited
        for.

        """
        # Condition - the mutex in this must be held when reading or writing
        # the member variables of Queue.
        #
        # Note - we might want to split this into a separate condition for each
        # set of tasks (ie, the waiting and finished queues), or even have a
        # condition for each task and just a mutex for the queues, to get
        # better scaling to more entries in the queues, and more threads
        # waiting on the tasks.
        self._condition = threading.Condition()

        # The next ID number to allocate to a request.
        self._next_id = 0

        # List, ordered ascending by priority, holding pairs of:
        #  (priority, list of requests waiting to begin)
        self._tasks_waiting = []

        # Map, keyed by request ID, holding requests which are in progress (ie,
        # have been returned by "get_next_request()", but haven't had results
        # returned to the queue).
        self._tasks_in_progress = {}

        # Map, keyed by request ID, holding the results of requests which have
        # finished.
        self._tasks_finished = {}

        # Flag, set to True when the queue is shutdown.
        self._shutdown = False

    def shutdown(self):
        """Close down the queue, discarding all requests.

        Any requests which have not yet been started are discarded; any
        requests which are in progress are converted to "noreturn" requests,
        and any requests which have finished but whose results have not been
        collected will have their results discarded.  Any callers waiting on
        the queue will be woken up, and have None returned (as if they had
        timed out).  Any future attempts to add requests to the queue will
        raise an exception.

        """
        self._condition.acquire()
        try:
            self._shutdown = True
            self._tasks_waiting = []
            self._tasks_in_progress.clear()
            self._tasks_finished = {}
            self._condition.notifyAll()
        finally:
            self._condition.release()

    def add(self, data, priority=0, noreturn=False, block=False):
        """Add a request to the queue.

        "data" is user-specified data describing the request.

        "priority" is the priority of the request.  This is an integer, such
        that higher numbers correspond to higher priorities - all requests of a
        given priority will be started before requests of a lower priority.

        Negative priorities are acceptable.  Also, storage and comparisons only
        consider those priority values which are specified, and no
        interpretation is put on priority values other than their ordering, so
        sparse priority values don't have an inordinate overhead (eg,
        performance should not suffer if you use "1000, 2000" for priority
        values instead of "1, 2").  This allows you to leave space between
        priority values for future expansion.

        Normally, this method returns immediately with an ID number for the
        request, which can be used to wait for its results.  ID numbers are
        unique only to requests currently in progress - they may be reused once
        the results of a request have been picked up.
        
        If "block" is True, the call will block until the results of the
        request are available.

        If "noreturn" is True, this method returns None, and any results from
        the request will be discarded.  Even if "block" is True, if "noreturn"
        is True the method will return immediately.

        """
        self._condition.acquire()
        try:
            if self._shutdown:
                raise ShutdownException("This queue has been shut down")

            # Make a Request object to hold the request, and allocate an ID to
            # it if needed.
            if noreturn:
                newid = None
            else:
                newid = self._next_id
                self._next_id += 1
            request = Request(data, priority, noreturn, newid, self)

            # Find the appropriate queue to append the request to, or make a
            # new queue if there are none for the specified priority.
            queuenum = bisect_left(self._tasks_waiting, priority)
            if queuenum == len(self._tasks_waiting):
                self._tasks_waiting.append((priority, [request]))
            else:
                queuepriority, queue = self._tasks_waiting[queuenum]
                if queuepriority == priority:
                    queue.append(request)
                else:
                    self._tasks_waiting.insert(queuenum, (priority, [request]))

            # Wake up anything waiting for a request to handle.
            self._condition.notifyAll()
            return newid
        finally:
            self._condition.release()

    def _finished(self, requestid, result):
        """Notify the queue that a request has finished.

        This should not be called for requests with noreturn set.

        If the request is not in the list of requests which are in progress
        (ie, it hadn't started yet, or hadn't finished yet, or the ID number is
        simply invalid), this will raise KeyError.

        """
        assert(requestid != None)
        self._condition.acquire()
        try:
            # del from a dict raises KeyError if the key isn't present.
            del self._tasks_in_progress[requestid]
            self._tasks_finished[requestid] = result

            # Wake up anything waiting for a request to finish.
            self._condition.notifyAll()
        finally:
            self._condition.release()

    def get_next_request(self, block=False, timeout=None, min_priority=None):
        """Get the next request to process.

        Returns a Request object holding the request.

        If "block" is True, the call will block until a request is available.
        Otherwise, the call will return immediately, and return None if no
        request is available.

        If "timeout" is specified, and "block" is True, the call will timeout
        after having waited for "timeout" seconds without a request becoming
        available.  After the timeout, the call will return None.  If "block"
        is not specified, "timeout" has no effect.

        Even if "block" is True and no timeout is specified, the call will
        return None if the queue is shutdown().

        If "min_priority" is specified, only requests with a priority of at
        least "min_priority" will be returned.  If there are no such requests,
        the call will block, or return None, just as it would there were no
        requests at all.

        """
        if timeout is None:
            endtime = None
        else:
            endtime = time.time() + timeout
        self._condition.acquire()
        try:
            while not self._shutdown:
                if len(self._tasks_waiting) != 0:
                    priority, waiting = self._tasks_waiting[-1]
                    if min_priority is None or priority >= min_priority:
                        request = waiting.pop(0)
                        assert(request.priority == priority)
                        if len(waiting) == 0:
                            del self.waiting[-1]
                        # Found a request to return.  Unless noreturn is set
                        # for it, add it to the _tasks_in_progress queue.
                        if not request.noreturn:
                            self._tasks_in_progress[request.id] = request
                        return request
                if not block:
                    return None
                if _wait_until_timeout(self._condition, endtime):
                    return None
            return None
        finally:
            self._condition.release()

    def get_finished(self, requestids, block=False, timeout=None):
        """Check to see if a request has been completed.

        "requestids" is either a single request ID, or a sequence of IDs.

        This call checks if any of the requests indicated in requestids have
        finished.  If they have, the result of the requests will be returned
        (not reliably the first in the list, though the present implementation
        will often return the earliest which has finished), together with the
        request ID for the request which is returned, as a RequestResult
        object.

        If "block" is True, the call will block until a request is available.
        Otherwise, the call will return immediately, and return None if no
        request is available.

        If "timeout" is specified, and "block" is True, the call will timeout
        after having waited for "timeout" seconds without a request becoming
        available.  After the timeout, the call will return None.  If "block"
        is not specified, "timeout" has no effect.

        Even if "block" is True and no timeout is specified, the call will
        return None if the queue is shutdown().

        """
        if timeout is None:
            endtime = None
        else:
            endtime = time.time() + timeout
        if isinstance(requestids, int):
            requestids = (requestids,)
        self._condition.acquire()
        try:
            while not self._shutdown:
                for requestid in requestids:
                    if requestid in self._tasks_finished:
                        result = self._tasks_finished.get(requestid)
                        return RequestResult(requestid, result)
                if not block:
                    return None
                if _wait_until_timeout(self._condition, endtime):
                    return None
            return None
        finally:
            self._condition.release()
